#!/usr/bin/python
# -*- coding: utf-8 -*-
# @Time: 20-8-4 下午7:59
# @Author: HuangChuQin
# @Desc: 保存指定股票的历史日数据，数据仓库名stock_k_data_daily_sh|sz.六位股票代码，如：stock_k_data_daily_sh.000001
# @File: stock_hist_data.py
# @Software: PyCharm

import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
import baostock as bs
import pymongo
import getopt
import time
import common.util_mysql
import common.utils

# 配置文件
profile = 'profile.ini'
# 日志对象
logger = None
# 配置对象
config = None

# 数据库参数
db_param = {}
# 数据库连接
db_conn = None

# mongodb集合
collections = None
# K线指标字典
field_dict = {}
# 股票代码
stock_code = None
stock_append = True
stock_init_all = False
stock_auto_run = False


def usage():
    print('********', sys.argv[0], '********')
    print('|--保存指定股票的历史日数据到数据仓库')
    print('|--数据仓库名stock_k_data_daily_sh|sz.六位股票代码，如：stock_k_data_daily_sh.000001')
    print('|--需要进入脚本所在的目录')
    print('|--', os.path.dirname(os.path.realpath(__file__)))
    print('|--参数说明')
    print('   -c/--code    接收股票代码，如：sh.000001')
    print('   -h/--help    帮助信息')
    print('   --auto-run   自动执行，追加A股当天的所有股票数据')
    print('   --init       初始化会把该股票的所有数据都爬取，默认只追加当天')
    print('   --init-all   初始化A股所有股票数据，耗时较长')
    print('********', sys.argv[0], '********')


def get_opt():
    global stock_code
    global stock_append
    global stock_init_all
    global stock_auto_run

    opts, args = getopt.getopt(sys.argv[1:], "hc:", ["help", "code", "init", "init-all", "auto-run"])
    for op, value in opts:
        if op in ("-c", "--code"):
            stock_code = value
            print("get_opt stock_code ", stock_code)
        elif op == "--init":
            stock_append = False
            print("get_opt stock_append true")
        elif op == "--init-all":
            stock_init_all = True
            print("get_opt stock_init_all true")
        elif op == "--auto-run":
            stock_auto_run = True
            print("get_opt stock_auto_run true")
        else:
            usage()
            sys.exit()
    if stock_code is None and stock_init_all is False and stock_auto_run is False:
        print('错误的执行方式，请指定股票代码')
        usage()
        sys.exit()


def my_init_conf():
    global profile
    global logger
    global config
    global db_param
    global collections

    # init confFile
    print("init config file:\t" + profile)
    config = common.utils.init_config(profile)
    if config is None:
        print("init config err, exit")
        sys.exit()
    print("init profile success")

    # init log
    log_config_file = config.get("COMMON", "LOG_CONFIG_PATH")
    log_config_name = config.get("COMMON", "LOG_NAME")
    print("log config : ", log_config_file, log_config_name)
    logger = common.utils.init_logger(log_config_file, log_config_name)
    if logger is None:
        print("init logger err, exit")
        sys.exit()
    print("init logger success")

    # init mysql
    db_param['SERVER'] = config.get("MYSQL", "SERVER")
    db_param['USER'] = config.get("MYSQL", "USER")
    db_param['PWD'] = config.get("MYSQL", "PWD")
    db_param['DBNAME'] = config.get("MYSQL", "DBNAME")

    client = pymongo.MongoClient('mongodb://127.0.0.1/')
    print(client.list_database_names())
    collections = client["stock_warehouse"]
    print(len(collections.list_collection_names()))


def db_conn_create():
    global db_conn
    db_conn = common.util_mysql.db_connect(db_param)
    if db_conn is None:
        logger.error('ConnDbMysql err')
        sys.exit()
    logger.info('connect mysql ok')


def db_conn_destroy():
    common.util_mysql.db_disconnect(db_conn)
    logger.info("destroy db_conn ok")


def get_stock_index():
    global field_dict
    (resultSet, rowCount, columnName) = common.util_mysql.fetch_multi_row(
        db_conn, "SELECT index_name FROM stock_base_info.stock_index_info_daily where index_platform='baostock'"
    )
    for result in resultSet:
        (index_name,) = result
        field_dict[index_name] = None


def main(code):
    print("######## 开始 " + str(time.time()))
    get_stock_index()
    from job_baostock.stock_data_append import StockDataAppend
    bs.login()
    # 初始化所有A股的所有数据
    if stock_init_all is True:
        collection_names = collections.list_collection_names()
        rs = bs.query_stock_basic()
        for stock in rs.data:
            if stock[5] is '0':
                print('***** 已退市，忽略不计 *****')
                print('股票信息：' + str(stock))
                continue
            if stock[4] is '2':
                print('***** 股指，忽略不计 *****')
                print('股票信息：' + str(stock))
                continue
            if 'stock_k_data_daily_'+stock[0] in collection_names:
                print('%s 已经在数据仓库中，现在追加当天数据' % stock[0])
                StockDataAppend(
                    mg_folder=collections, bao_stock=bs, stock_code=stock[0], field_dict=field_dict,
                )
                continue
            StockDataAppend(
                mg_folder=collections, bao_stock=bs, stock_code=stock[0], field_dict=field_dict,
                is_init=True, date=stock[2]
            )
    # 追加所有A股当天的数据
    elif stock_auto_run is True:
        rs = bs.query_stock_basic()
        for stock in rs.data:
            if stock[5] is '0':
                print('***** 已退市，忽略不计 *****')
                print('股票信息：' + str(stock))
                continue
            if stock[4] is '2':
                print('***** 股指，忽略不计 *****')
                print('股票信息：' + str(stock))
                continue
            StockDataAppend(
                mg_folder=collections, bao_stock=bs, stock_code=stock[0], field_dict=field_dict,
            )
    # 单只股票操作
    else:
        StockDataAppend(
            mg_folder=collections, bao_stock=bs, stock_code=code, field_dict=field_dict, is_init=stock_append is False
        )
    bs.logout()
    print("######## 结束 " + str(time.time()))


if __name__ == "__main__":
    get_opt()

    my_init_conf()
    db_conn_create()
    try:
        # main('sh.000001')
        main(stock_code)
    except:
        print("Unexpected error:", sys.exc_info()[0])

    db_conn_destroy()

    # 退出程序
    logger.info("run end, exit")
